from funasr import AutoModel

chunk_size = [0, 5, 3] #[0, 10, 5] 600ms, [0, 8, 4] 480ms
encoder_chunk_look_back = 4 #number of chunks to lookback for encoder self-attention
decoder_chunk_look_back = 1 #number of encoder chunks to lookback for decoder cross-attention

model = AutoModel(model="paraformer-zh-streaming", disable_update=True)

import soundfile
import os
import pyaudio
import numpy as np
import wave

# 音频录制参数
CHUNK = 4800
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000

p = pyaudio.PyAudio()

# 获取设备信息
info = p.get_host_api_info_by_index(0)
num_devices = info.get('deviceCount')

# 打印设备信息
for i in range(num_devices):
    device_info = p.get_device_info_by_index(i)
    print(f"Device index: {i}, Device name: {device_info.get('name')}")

stream = p.open(format=FORMAT,
                channels=CHANNELS,
                rate=RATE,
                input=True,
                input_device_index=3,  # 指定录音设备
                frames_per_buffer=CHUNK)

cache = {}
print("start to read audio")
index = 0
frames = []
while index < 500:
    input_data = stream.read(CHUNK, exception_on_overflow=False)
    audio_chunk = np.frombuffer(input_data, dtype=np.int16)
    # audio_chunk = input_data
    print(f"len = {len(audio_chunk)}")
    is_final = 0
    res = model.generate(input=audio_chunk, cache=cache, is_final=is_final, chunk_size=chunk_size, encoder_chunk_look_back=encoder_chunk_look_back, decoder_chunk_look_back=decoder_chunk_look_back)
    #res = model.generate(input=data)
    print(res)
    index = index + 1

print("all done")
# wf = wave.open("test-record.wav", 'wb')
# wf.setnchannels(CHANNELS)
# wf.setsampwidth(p.get_sample_size(FORMAT))
# wf.setframerate(RATE)
# wf.writeframes(b''.join(frames))
# wf.close()
